{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'\n",
    "\n",
    "import sys \n",
    "sys.path.append(\"../seg/\") \n",
    "\n",
    "from utils.data_generator import ImageDataGenerator\n",
    "from utils.helpers import get_dataset_info, check_related_path\n",
    "from utils.callbacks import LearningRateScheduler\n",
    "from utils.optimizers import *\n",
    "from utils.losses import *\n",
    "from utils.learning_rate import *\n",
    "from utils.metrics import *\n",
    "from utils import utils\n",
    "from builders import builder\n",
    "import tensorflow as tf\n",
    "import argparse\n",
    "import numpy as np            \n",
    "np.random.seed(0)\n",
    "tf.compat.v1.set_random_seed(0)\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "\n",
    "\n",
    "# check related paths\n",
    "paths = check_related_path(os.getcwd())\n",
    "\n",
    "# get image and label file names for training and validation\n",
    "train_image_names, train_label_names, valid_image_names, valid_label_names, test_image_names,test_label_names = get_dataset_info('../dataset/segmentation')\n",
    "\n",
    "# build the model\n",
    "net, base_model = builder(2, (256, 256), 'DeepLabV3Plus', 'Xception-DeepLab')\n",
    " \n",
    "net.load_weights('../checkpoint/seg.h5')\n",
    "\n",
    "loss = {\n",
    "    'tf.concat':metric_loss('c5',50,0.01),\n",
    "    'seg':focal_loss() \n",
    "}\n",
    "\n",
    "# compile the model\n",
    "net.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=3e-4),\n",
    "        loss=loss,\n",
    "        metrics={'seg':[BiIOU()]},\n",
    "        run_eagerly=True,\n",
    "        loss_weights={'tf.concat':0.01,'seg': 1}#tf_op_layer_feature\n",
    ")\n",
    "\n",
    "valid_gen = ImageDataGenerator()\n",
    "\n",
    "valid_generator = valid_gen.flow(images_list=valid_image_names,\n",
    "                                 labels_list=valid_label_names,\n",
    "                                 num_classes=2,\n",
    "                                 batch_size=4,\n",
    "                                 target_size=(256, 256))\n",
    "net.evaluate(valid_generator,\n",
    "                  batch_size=4,\n",
    "                  max_queue_size=10,\n",
    "                  workers=os.cpu_count(),\n",
    "                  use_multiprocessing=False,\n",
    "                 steps = 59//4\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# classifier - CNN\n",
    "\n",
    "import sys \n",
    "sys.path.append(\"../class/\")\n",
    "import io\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tensorflow_addons as tfa\n",
    "import tensorflow.keras as keras\n",
    "from sklearn.cluster import KMeans\n",
    "\n",
    "origin_resnet = keras.applications.ResNet50V2(weights='imagenet')\n",
    "resnet = tf.keras.Model(inputs=origin_resnet.input,outputs=origin_resnet.layers[-2].output)\n",
    "\n",
    "class_model = tf.keras.Sequential([\n",
    "    resnet,\n",
    "    tf.keras.layers.Lambda(lambda x: tf.math.l2_normalize(x, axis=1)) # L2 normalize embeddings\n",
    "])\n",
    "\n",
    "class_model.load_weights('../checkpoint/class.h5')\n",
    "\n",
    "class_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=3e-4),\n",
    "        loss=tfa.losses.TripletSemiHardLoss(margin=0.01),\n",
    "        run_eagerly=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# classifier - KNN\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "train_dir = '../dataset/classification/train'\n",
    "test_dir = '../dataset/classification/valid'\n",
    "batch_size = 64\n",
    "crop_size = 224\n",
    "\n",
    "train_dataset = tf.keras.preprocessing.image_dataset_from_directory(\n",
    "    train_dir,\n",
    "    seed=1,\n",
    "    image_size=(crop_size, crop_size),\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    ")\n",
    "\n",
    "def _normalize_img(img, label):\n",
    "    img = tf.keras.applications.resnet_v2.preprocess_input(img)\n",
    "    return (img, label)\n",
    "\n",
    "train_dataset = train_dataset.map(_normalize_img)\n",
    "\n",
    "def get_dataset(dataset):\n",
    "    X=tf.zeros([0,2048])\n",
    "    Y=tf.zeros([0],dtype=tf.int32)\n",
    "    for sample in dataset:\n",
    "        x = class_model.predict(sample[0])\n",
    "        y = sample[1]\n",
    "        X = tf.concat([X,x],0)\n",
    "        Y = tf.concat([Y,y],0)\n",
    "    return X,Y\n",
    "\n",
    "# KNN\n",
    "\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "\n",
    "X_train,Y_train = get_dataset(train_dataset)\n",
    "\n",
    "\n",
    "neigh = KNeighborsClassifier(n_neighbors=1)\n",
    "neigh.fit(X_train, Y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Extractor\n",
    "\n",
    "import sys \n",
    "sys.path.append(\"../seg/\") \n",
    "from utils.helpers import check_related_path, get_colored_info, color_encode\n",
    "from utils.utils import load_image,load_label, decode_one_hot\n",
    "from keras_applications import imagenet_utils\n",
    "from builders import builder\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import argparse\n",
    "import sys\n",
    "import cv2\n",
    "import os\n",
    "from skimage import measure\n",
    "from utils.metrics import *\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def compute_iou(x,y):\n",
    "    set_x = set()\n",
    "    set_y = set()\n",
    "    for coord in x:\n",
    "        set_x.add(coord[0]*10000000+coord[1])\n",
    "    for coord in y:\n",
    "        set_y.add(coord[0]*10000000+coord[1])\n",
    "    return len(set_x & set_y) / len(set_x | set_y)\n",
    "\n",
    "def get_connected_areas(img,seg, margin=2):\n",
    "    cnt = lambda x,a:(x==a).sum()\n",
    "    label, num = measure.label(seg, background=None, return_num=True, connectivity=2)\n",
    "    properties = measure.regionprops(label)\n",
    "    for prop in properties:\n",
    "        if prop.area < 20: \n",
    "            properties.remove(prop)\n",
    "    \n",
    "    images = []\n",
    "    labels = []\n",
    "    coords = []\n",
    "    for i,prop in enumerate(properties):\n",
    "        \n",
    "        coord = prop.coords\n",
    "        bbox = prop.bbox\n",
    "        len1, len2 = bbox[2]-bbox[0], bbox[3]-bbox[1]\n",
    "        \n",
    "        if len1<20 and len2<20:\n",
    "            continue\n",
    "        \n",
    "        b1 = max(0, bbox[0]-int(len1/margin) )\n",
    "        b2 = bbox[2]+int(len1/margin)\n",
    "        b3 = max(0, bbox[1]-int(len2/margin))\n",
    "        b4 = bbox[3]+int(len2/margin)\n",
    "\n",
    "        sub_image = img[b1 : b2 , b3 : b4]\n",
    "        sub_seg = seg[b1 : b2 , b3 : b4]\n",
    "        \n",
    "        label = 0 if cnt(sub_seg,1)>cnt(sub_seg,2) else 1\n",
    "        \n",
    "        coords.append(coord)\n",
    "        images.append(sub_image)\n",
    "        labels.append(label)\n",
    "        \n",
    "    \n",
    "    return images,labels,coords\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
